import rclpy
from rclpy.node import Node
from mit_msgs.msg import MITJointCommands, MITLowState
from sensor_msgs.msg import JointState, Imu
from unitree_hg.msg import LowCmd, LowState
from .crc import CRC

from rclpy.executors import MultiThreadedExecutor
from rclpy.callback_groups import MutuallyExclusiveCallbackGroup

G1_29DOF_JOINT_NAME = [
    "l_leg_hip_pitch",
    "l_leg_hip_roll",
    "l_leg_hip_yaw",
    "l_leg_knee",
    "l_leg_ankle_pitch",
    "l_leg_ankle_roll",
    "r_leg_hip_pitch",
    "r_leg_hip_roll",
    "r_leg_hip_yaw",
    "r_leg_knee",
    "r_leg_ankle_pitch",
    "r_leg_ankle_roll",
    "waist_yaw",
    "waist_roll",
    "waist_pitch",
    "l_shoulder_pitch",
    "l_shoulder_roll",
    "l_shoulder_yaw",
    "l_elbow",
    "l_wrist_roll",
    "l_wrist_pitch",
    "l_wrist_yaw",
    "r_shoulder_pitch",
    "r_shoulder_roll",
    "r_shoulder_yaw",
    "r_elbow",
    "r_wrist_roll",
    "r_wrist_pitch",
    "r_wrist_yaw"
]

DEFAULT_JOINT_POS = [
    -0.20, 0.00, 0.00, 0.42, -0.23, 0.00, 
    -0.20, 0.00, 0.00, 0.42, -0.23, 0.00,
    0.00, 0.00, 0.00, 
    0.35, 0.16, 0.00, 0.87, 0.00, 0.00, 0.00,
    0.35, -0.16, 0.00, 0.87, 0.00, 0.00, 0.00
]


class TopicConvert(Node):

    def __init__(self):
        super().__init__('topic_convert_node')

        # G1 State Handler
        g1_state_callback_group = MutuallyExclusiveCallbackGroup()
        self.g1_lowstate_sub = self.create_subscription(LowState, '/lowstate', self.G1_lowstate_callback, 10, callback_group=g1_state_callback_group)
        self.g1_mitstate_pub = self.create_publisher(MITLowState, '/human_lower_state', 1, callback_group=g1_state_callback_group)
        # 队列长度设为1，仅接受最新。若还不可以，就加允许丢包

        # MIT Command Handler
        mit_cmd_callback_group = MutuallyExclusiveCallbackGroup()
        self.mit_cmd_sub = self.create_subscription(MITJointCommands, '/human_lower_command', self.Mit_command_callback, 10, callback_group=mit_cmd_callback_group)
        self.converted_mit_pub = self.create_publisher(LowCmd, '/lowcmd', 10, callback_group=mit_cmd_callback_group)

        self.cmd = LowCmd()

        self.crc = CRC()  # Initialize the CRC class

        # self.timer = self.create_timer(0.1, self.timer_callback)

        self.mode_machine = 2  # 默认为29自由度，避免启动顺序错误导致异常
        self.wireless_remote = None
        # self.idle_mode = True  # 是否进入阻尼模式 
        self.last_idle_button = 0
    
    def G1_lowstate_callback(self, msg:LowState):
        """
        接收底层传上来的状态
        转换为MIT消息发送出去
        同时处理手柄按键
        """
        # 判断是否进入阻尼模式
        # self.wireless_remote = msg.wireless_remote
        # if not self.idle_mode and self.wireless_remote[2] == 0 and self.last_idle_button == 1:
        #     self.idle_mode = True
        #     self.DampingMode()
        #     self.get_logger().info('进入阻尼模式')
        # elif self.idle_mode and self.wireless_remote[2] == 0 and self.last_idle_button == 1:
        #     self.idle_mode = False
        #     self.get_logger().info('进入运动模式')

        # self.last_idle_button = self.wireless_remote[2]

        # print(f'wireless_remote: {self.wireless_remote}')

        self.mode_machine = msg.mode_machine

        self.g1_mit_low_state = MITLowState()

        # joint_state 转换
        joint_state_msg = JointState()
        joint_state_msg.name = [G1_29DOF_JOINT_NAME[i] for i in range(len(msg.motor_state[:29]))]  # 前29有值，后为冗余位，全0
        joint_state_msg.position = [m.q for m in msg.motor_state]
        joint_state_msg.velocity = [m.dq for m in msg.motor_state]
        joint_state_msg.effort = [m.tau_est for m in msg.motor_state]  # TODO: 是对应的tau_est吗？
        self.g1_mit_low_state.joint_states = joint_state_msg

        # imu_state 转换
        imu_msg = Imu()
        imu_msg.orientation.x = float(msg.imu_state.quaternion[1])
        imu_msg.orientation.y = float(msg.imu_state.quaternion[2])
        imu_msg.orientation.z = float(msg.imu_state.quaternion[3])
        imu_msg.orientation.w = float(msg.imu_state.quaternion[0])
        imu_msg.angular_velocity.x = float(msg.imu_state.gyroscope[0])
        imu_msg.angular_velocity.y = float(msg.imu_state.gyroscope[1])
        imu_msg.angular_velocity.z = float(msg.imu_state.gyroscope[2])
        imu_msg.linear_acceleration.x = float(msg.imu_state.accelerometer[0])
        imu_msg.linear_acceleration.y = float(msg.imu_state.accelerometer[1])
        imu_msg.linear_acceleration.z = float(msg.imu_state.accelerometer[2])
        self.g1_mit_low_state.imu = imu_msg

        self.g1_mitstate_pub.publish(self.g1_mit_low_state)
        
    def Mit_command_callback(self, mit_cmds:MITJointCommands):  
        self.cmd.mode_pr = 0  # 脚踝解算模式
        self.cmd.mode_machine = self.mode_machine
        for i, mit_cmd in enumerate(mit_cmds.commands):
            if i < len(self.cmd.motor_cmd):
                self.get_logger().debug(f'电机个数：{len(self.cmd.motor_cmd)}')
                self.cmd.motor_cmd[i].mode = 1  # 1: enable, 0: disable
                self.cmd.motor_cmd[i].tau = mit_cmd.eff
                self.cmd.motor_cmd[i].q = mit_cmd.pos  # target
                self.cmd.motor_cmd[i].dq = mit_cmd.vel  # velocity
                self.cmd.motor_cmd[i].kp = mit_cmd.kp   
                self.cmd.motor_cmd[i].kd = mit_cmd.kd
        
        self.cmd.crc = self.crc.Crc(self.cmd)   
        self.converted_mit_pub.publish(self.cmd)
        # if not self.idle_mode:  # 不在阻尼模式下才发布
        #     self.converted_mit_pub.publish(self.cmd)
        # else:
        #     self.DampingMode()

    # def timer_callback(self):
        # if self.idle_mode:
        #     self.DampingMode()

    # def DampingMode(self):
    #     pass
        # self.cmd.mode_pr = 0  # 脚踝解算模式
        # self.cmd.mode_machine = self.mode_machine
        # for i in range(len(G1_29DOF_JOINT_NAME)):
        #     self.cmd.motor_cmd[i].mode = 1  # 1: enable, 0: disable
        #     self.cmd.motor_cmd[i].tau = 0.0
        #     self.cmd.motor_cmd[i].q = DEFAULT_JOINT_POS[i]  # target
        #     self.cmd.motor_cmd[i].dq = 0.0  # velocity
        #     self.cmd.motor_cmd[i].kp = 50.0  
        #     self.cmd.motor_cmd[i].kd = 1.0
        
        # self.cmd.crc = self.crc.Crc(self.cmd)   
        # self.converted_mit_pub.publish(self.cmd)

def main(args=None):
    rclpy.init(args=args)

    mit_g1_controller = TopicConvert()
    mit_g1_controller.get_logger().info('topic converter is running!')
    
    executor = MultiThreadedExecutor(num_threads=4)
    executor.add_node(mit_g1_controller)
    executor.spin()

    # mit_g1_controller.destroy_node()
    rclpy.shutdown()


if __name__ == '__main__':
    main()
